from functools import wraps
import inspect
import time
import grpc
from typing import Callable, Optional
from enum import IntEnum

from common.utils import logger

# grpc错误码，需注意与go中保持一致
class ErrorCode(IntEnum):
    SERVER_COMMON_ERROR = 100001
    REQUEST_PARAM_ERROR = 100002
    TOKEN_EXPIRE_ERROR = 100003  
    TOKEN_GENERATE_ERROR = 100004
    DB_ERROR = 100005
    DB_UPDATE_AFFECTED_ZERO_ERROR = 100006

_error_messages = {
    ErrorCode.SERVER_COMMON_ERROR: "服务器开小差了，稍后再来试一试",
    ErrorCode.REQUEST_PARAM_ERROR: "参数错误",
    ErrorCode.TOKEN_EXPIRE_ERROR: "token失效，请重新登录",
    ErrorCode.TOKEN_GENERATE_ERROR: "生成token失败",
    ErrorCode.DB_ERROR: "数据库繁忙，请稍后重试",
    ErrorCode.DB_UPDATE_AFFECTED_ZERO_ERROR: "更新数据库影响行数为0"
}

class MyException(Exception):
    def __init__(self, error_code: ErrorCode, error_log: str, *, detail: str=None):
        self.error_code = error_code
        self.detail = detail if detail is not None else _error_messages[error_code]
        super().__init__(f"MyException: [{error_code}] {error_log}: {self.detail}")

def _default_handler(e: Exception, context: grpc.ServicerContext, response_type: type, call_info: dict = None):
    error_log = ""
    if isinstance(e, MyException):
        error_log = str(e)
        context.set_code(e.error_code)
        context.set_details(e.detail)
    elif isinstance(e, grpc.RpcError):
        context.set_code(e.code())
        error_log = f"RPC error: {e.details()}"
        context.set_details(e.details())
    else:
        context.set_code(grpc.StatusCode.INTERNAL)
        error_log = f"Internal error: {str(e)}"
        context.set_details(error_log)
    
    logger.my_logger().error(f'{call_info}: {error_log}')
    
    return response_type() if response_type else None

def with_span_handler():
    """
    支持trace注入处理的gRPC装饰器
    """
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(self, request, context):
            tracing_util = getattr(self, 'tracing_util', None)
            
            with tracing_util(context) as span:
                result = func(self, request, context)
                if inspect.isgenerator(result):
                    def wrapped_gen():
                        with tracing_util(context) as span:
                            for item in result:
                                yield item
                    return wrapped_gen()
                else:
                    if span:
                        span.end()
                    return result
            
        return wrapper
    return decorator

def grpc_exception_handler(
    response_type: type = None,
    exception_handler: Optional[Callable] = _default_handler
):
    """
    支持自定义异常处理的gRPC装饰器
    :param response_type: 返回的protobuf消息类型
    :param exception_handler: 自定义异常处理方法
    """
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(self, request, context):
            call_info = {
                'module': func.__module__,
                'class': self.__class__.__name__ if hasattr(self, '__class__') else None,
                'function': func.__name__,
            }
            
            try:
                result = func(self, request, context)
                if inspect.isgenerator(result):
                    try:
                        yield from result
                    except Exception as e:
                        yield exception_handler(e, context, response_type, call_info)
                else:
                    return result
            except Exception as e:
                return exception_handler(e, context, response_type, call_info)

        return wrapper
    return decorator

def log_time_handler():
    """
    支持记录处理时间日志的gRPC装饰器
    """
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            call_info = {
                'module': func.__module__,
                'class': self.__class__.__name__ if hasattr(self, '__class__') else None,
                'function': func.__name__,
            }
            
            start = time.time()
            result = func(self, *args, **kwargs)
            if inspect.isgenerator(result):
                def wrapped_gen():
                    for item in result:
                        yield item
                    duration_ms = round((time.time() - start) * 1000, 1)
                    logger.my_logger().info(f'process {call_info}', extra={'duration': duration_ms})
                return wrapped_gen()
            else:
                duration_ms = round((time.time() - start) * 1000, 1)
                logger.my_logger().info(f'process {call_info}', extra={'duration': duration_ms})
                return result
        
        return wrapper
    return decorator
            
def my_grpc_handler(
    response_type=None,
    exception_handler=_default_handler,
    enable_trace=True,
    enable_log=True
):
    def decorator(func):
        if enable_log:
            func = log_time_handler()(func)
        if response_type is not None:
            func = grpc_exception_handler(response_type, exception_handler)(func)
        if enable_trace:
            func = with_span_handler()(func)
        return func
    return decorator
